from flask import Flask, request, jsonify
import requests
import json

API_KEY = "x5rtRCxfj4XI20UrhrAVUbwU"
SECRET_KEY = "ochldfQI2S9iJFJij7nfW92v3F20kMm8"

app = Flask(__name__)

@app.route('/chat_with_ernie', methods=['POST'])
def chat_with_ernie():

    # 获取access_token
    access_token = get_access_token()

    url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-speed-128k?access_token={access_token}"

    # 从请求体中获取payload
    payload_obj = request.get_json(force=True)

    payload = json.dumps(payload_obj)

    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)

    print(response.text)

    # 返回结果给前端

    return jsonify(response.json()), response.status_code


def get_access_token():
    """
    使用 AK，SK 生成鉴权签名（Access Token）
    :return: access_token，或是None(如果错误)
    """
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {"grant_type": "client_credentials", "client_id": API_KEY, "client_secret": SECRET_KEY}
    return str(requests.post(url, params=params).json().get("access_token"))

if __name__ == '__main__':
    app.run(debug=True)
